import numpy as np
import cv2
from scipy.special import softmax
from typing import Dict, List, Optional, Union, Tuple

from .image_preprocessor import ImagePreprocessor
from .model_postprocessor import ModelPostprocessor

class YOLOPreprocessor(ImagePreprocessor):
    """YOLO风格的图像预处理"""
    
    def __init__(self, target_width: int = 928, target_height: int = 512, pad_value: Tuple = (114, 114, 114)):
        self.target_width = target_width
        self.target_height = target_height
        self.pad_value = pad_value
    
    def resize_and_pad(self, image: np.ndarray) -> Tuple[np.ndarray, float, float]:
        original_height, original_width = image.shape[:2]
        ratio_w = 1.0
        ratio_h = 1.0

        if original_width > self.target_width or original_height > self.target_height:
            ratio_w = self.target_width / original_width
            ratio_h = self.target_height / original_height
            ratio = min(ratio_w, ratio_h)

            new_width = int(original_width * ratio)
            new_height = int(original_height * ratio)
            image_resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
        else:
            image_resized = image.copy()

        resized_height, resized_width = image_resized.shape[:2]
        delta_w = self.target_width - resized_width
        delta_h = self.target_height - resized_height

        if delta_w > 0 or delta_h > 0:
            image_padded = cv2.copyMakeBorder(
                image_resized,
                top=0, bottom=delta_h, left=0, right=delta_w,
                borderType=cv2.BORDER_CONSTANT,
                value=self.pad_value
            )
        else:
            image_padded = image_resized

        return image_padded, ratio_w, ratio_h

    def preprocess(self, image: np.ndarray) -> Tuple[np.ndarray, Dict]:
        """
        预处理图像为YOLO模型输入格式。

        Args:
            image (np.ndarray): 输入图像，形状为 (H, W, 3)。

        Returns:
            Tuple[np.ndarray, Dict]: 包含预处理后的图像张量和元数据的元组。
        """
        # 调整图像大小并填充
        resized_padded_image, ratio_w, ratio_h = self.resize_and_pad(image)
        # 归一化
        image_normalized = resized_padded_image / 255.0
        # 转换为CHW格式
        image_transposed = np.transpose(image_normalized, (2, 0, 1))
        # 添加批次维度
        input_tensor = np.expand_dims(image_transposed, axis=0).astype(np.float32)
        
        meta = {
            "resized_padded_image": resized_padded_image,
            "ratio_w": ratio_w,
            "ratio_h": ratio_h,
            "original_shape": image.shape,
            "target_shape": (self.target_width, self.target_height)
        }
        
        return input_tensor, meta


class YOLOPostprocessor(ModelPostprocessor):
    def __init__(self, confidence_threshold: float = 0.25, iou_threshold: float = 0.45):
        self.confidence_threshold = confidence_threshold
        self.iou_threshold = iou_threshold

    def postprocess(self, model_output, meta: Dict) -> Dict:
        output = model_output[0][0]  # (5, 17661)
        output = np.transpose(output, (1, 0))  # (17661, 5)

        # 过滤低置信度框
        valid_boxes = output[output[:, 4] > self.confidence_threshold]
        
        if len(valid_boxes) == 0:
            return {"defect_boxes": []}

        # 提取框坐标和置信度
        boxes = valid_boxes[:, :4]  # xc, yc, w, h
        scores = valid_boxes[:, 4]  # objectness 置信度

        # 执行 NMS
        keep_indices = self._nms(boxes, scores, self.iou_threshold)
        final_boxes = boxes[keep_indices]
        final_scores = scores[keep_indices]

        # 转换为 x1y1x2y2 格式
        defect_boxes = []
        for i in range(len(final_boxes)):
            xc, yc, w, h = final_boxes[i]
            x1 = float(xc - w / 2)
            y1 = float(yc - h / 2)
            x2 = float(xc + w / 2)
            y2 = float(yc + h / 2)

            box_dict = {
                "x1": x1,
                "y1": y1,
                "x2": x2,
                "y2": y2,
                "confidence": float(final_scores[i]),
            }
            defect_boxes.append(box_dict)

        return {"defect_boxes": defect_boxes}

    def _nms(self, boxes, scores, iou_threshold):
        """
        NMS 实现，仅基于置信度，不考虑类别
        boxes: (N, 4) -> xc, yc, w, h
        scores: (N,)
        """
        if len(boxes) == 0:
            return []

        # 转换为 x1y1x2y2 格式
        x1 = boxes[:, 0] - boxes[:, 2] / 2
        y1 = boxes[:, 1] - boxes[:, 3] / 2
        x2 = boxes[:, 0] + boxes[:, 2] / 2
        y2 = boxes[:, 1] + boxes[:, 3] / 2

        areas = (x2 - x1) * (y2 - y1)
        order = scores.argsort()[::-1]

        keep = []
        while len(order) > 0:
            i = order[0]
            keep.append(i)

            # 计算交集坐标
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            # 交集面积
            w = np.maximum(0.0, xx2 - xx1)
            h = np.maximum(0.0, yy2 - yy1)
            inter = w * h

            # IoU
            ovr = inter / (areas[i] + areas[order[1:]] - inter)

            # 保留 IoU 小于阈值的框
            inds = np.where(ovr <= iou_threshold)[0]
            order = order[inds + 1]

        return keep


class YOLOPostprocessorWithClass_(ModelPostprocessor):
    def __init__(self, confidence_threshold=0.01, iou_threshold=0.45):
        self.confidence_threshold = confidence_threshold
        self.iou_threshold = iou_threshold
        self.class_names = ['quejiao', 'qipao', 'liewen', 'shayan']
        self.class_labels = ['缺角', '镀层起泡', '裂纹', '砂眼']
        # 类别专属 confidence 阈值
        self.class_conf_thresholds = {
            'quejiao': 0.15,   # 高精度，适当过滤低置信度
            'qipao': 0.015,    # 提升召回
            'liewen': 0.005,   # 极低样本、低召回，保留更多预测
            'shayan': 0.08     # 精度较低，适当过滤误检
        }

    def postprocess(self, model_output, meta: dict) -> dict:
        output = model_output[0]
        output = output[0]

        # 转置为 (26880, 9)
        output = np.transpose(output, (1, 0))  # (26880, 9)

        # 提取各部分
        boxes = output[:, :4]       # (26880, 4)
        obj_conf = output[:, 4]     # (26880,)
        class_conf = output[:, 5:]  # (26880, 4)

        # softmax 分类置信度
        class_conf = softmax(class_conf, axis=1)

        # 总置信度 = obj_conf * max_class_conf
        scores = obj_conf * np.max(class_conf, axis=1)  # (26880,)
        class_ids = np.argmax(class_conf, axis=1)       # (26880,)

        # 【第一步过滤】粗略 confidence 阈值过滤
        valid_indices = np.where(scores > self.confidence_threshold)[0]

        if len(valid_indices) == 0:
            return {"defect_boxes": [], "category": "normal"}

        valid_boxes = boxes[valid_indices]
        valid_scores = scores[valid_indices]
        valid_class_ids = class_ids[valid_indices]

        # 【第二步过滤】类别专属 confidence 阈值过滤
        final_valid = []
        for i in range(len(valid_indices)):
            cls_id = valid_class_ids[i]
            score = valid_scores[i]
            class_name = self.class_names[cls_id]

            # 获取该类别的专属阈值
            threshold = self.class_conf_thresholds.get(class_name, self.confidence_threshold)

            if score >= threshold:
                final_valid.append(i)  # 保留该预测框的索引

        # 更新 valid 数据
        valid_boxes = valid_boxes[final_valid]
        valid_scores = valid_scores[final_valid]
        valid_class_ids = valid_class_ids[final_valid]

        # 【第三步】NMS
        nms_indices = self._nms(valid_boxes, valid_scores, self.iou_threshold)
        final_boxes = valid_boxes[nms_indices]
        final_scores = valid_scores[nms_indices]
        final_class_ids = valid_class_ids[nms_indices]

        # 构造输出
        result = []
        for i in range(len(final_boxes)):
            xc, yc, w, h = final_boxes[i]
            x1 = float(xc - w / 2)
            y1 = float(yc - h / 2)
            x2 = float(xc + w / 2)
            y2 = float(yc + h / 2)
            categoryId = int(final_class_ids[i])

            result.append({
                "x1": x1,
                "y1": y1,
                "x2": x2,
                "y2": y2,
                "confidence": float(final_scores[i]),
                "categoryId": categoryId,
                "categoryName": self.class_names[categoryId],
                "categoryLabel": self.class_labels[categoryId],
            })

        # 有缺陷，不是 normal
        return {"defect_boxes": result, "category": "defect"}

    def _nms(self, boxes, scores, iou_threshold=0.5):
        x1 = boxes[:, 0] - boxes[:, 2] / 2
        y1 = boxes[:, 1] - boxes[:, 3] / 2
        x2 = boxes[:, 0] + boxes[:, 2] / 2
        y2 = boxes[:, 1] + boxes[:, 3] / 2

        areas = (x2 - x1) * (y2 - y1)
        order = scores.argsort()[::-1]

        keep = []
        while len(order) > 0:
            i = order[0]
            keep.append(i)

            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(0.0, xx2 - xx1)
            h = np.maximum(0.0, yy2 - yy1)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)

            inds = np.where(ovr <= iou_threshold)[0]
            order = order[inds + 1]

        return keep

class YOLOPostprocessorWithClass(ModelPostprocessor):
    def __init__(
        self,
        class_names: List[str] = ['quejiao', 'qipao', 'liewen', 'shayan'],
        class_labels: List[str] = ['缺角', '镀层起泡', '裂纹', '砂眼'],
        class_conf_thresholds: Dict[str, float] = {
            'quejiao': 0.15,   # 高精度，适当过滤低置信度
            'qipao': 0.015,    # 提升召回
            'liewen': 0.005,   # 极低样本、低召回，保留更多预测
            'shayan': 0.08     # 精度较低，适当过滤误检
        },
        default_conf_threshold: float = 0.1,
        iou_threshold: float = 0.5,
        image_size: Optional[tuple] = None
    ):
        """
        改进的YOLO后处理器，支持类别专属阈值和更健壮的处理
        
        参数:
            class_names: 类别名称列表
            class_labels: 类别显示标签列表
            class_conf_thresholds: 每个类别的专属置信度阈值
            default_conf_threshold: 默认置信度阈值
            iou_threshold: NMS的IoU阈值
            image_size: (width, height) 用于边界检查
        """
        self.class_names = class_names
        self.class_labels = class_labels
        self.class_conf_thresholds = class_conf_thresholds
        self.confidence_threshold = default_conf_threshold
        self.iou_threshold = iou_threshold
        self.image_size = image_size
        
        # 验证输入
        if len(class_names) != len(class_labels):
            raise ValueError("class_names and class_labels must have same length")
        
        for cls in class_conf_thresholds:
            if cls not in class_names:
                raise ValueError(f"Class {cls} in class_conf_thresholds not found in class_names")

    def postprocess(self, model_output: Union[List, np.ndarray], meta: dict) -> dict:
                # 输出解析 (适用于YOLOv11的ONNX输出格式)
        predictions = np.squeeze(model_output[0]).T  # [26880, 8]
        
        # 过滤低置信度检测
        scores = np.max(predictions[:, 4:8], axis=1)
        mask = scores >= self.confidence_threshold
        predictions = predictions[mask]
        scores = scores[mask]
        
        if len(predictions) == 0:
            return []
        
        # 获取类别和置信度
        class_ids = np.argmax(predictions[:, 4:8], axis=1)
        boxes = predictions[:, :4]

        indices = self._nms(boxes, scores, self.iou_threshold)
        
        # 构造输出
        result = []
        for i in range(len(indices)):
            xc, yc, w, h = boxes[i]
            x1 = float(xc - w / 2)
            y1 = float(yc - h / 2)
            x2 = float(xc + w / 2)
            y2 = float(yc + h / 2)
            categoryId = int(class_ids[i])
            confidence = float(scores[i])

            result.append({
                "x1": x1,
                "y1": y1,
                "x2": x2,
                "y2": y2,
                "confidence": confidence,
                "categoryId": categoryId,
                "categoryName": self.class_names[categoryId],
                "categoryLabel": self.class_labels[categoryId],
            })

        # 有缺陷，不是 normal
        return {"defect_boxes": result, "category": "defect"}

    def _nms(self, boxes, scores, iou_threshold=0.5):
        x1 = boxes[:, 0] - boxes[:, 2] / 2
        y1 = boxes[:, 1] - boxes[:, 3] / 2
        x2 = boxes[:, 0] + boxes[:, 2] / 2
        y2 = boxes[:, 1] + boxes[:, 3] / 2

        areas = (x2 - x1) * (y2 - y1)
        order = scores.argsort()[::-1]

        keep = []
        while len(order) > 0:
            i = order[0]
            keep.append(i)

            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(0.0, xx2 - xx1)
            h = np.maximum(0.0, yy2 - yy1)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)

            inds = np.where(ovr <= iou_threshold)[0]
            order = order[inds + 1]

        return keep
